use std::cmp::Ordering;
use std::mem;

/// Type of heap: min heap or max heap.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum HeapType {
    Min,
    Max,
}

/// A trait for comparing values in the heap.
pub trait HeapCompare<T> {
    fn compare(&self, a: &T, b: &T) -> Ordering;
}

/// A binomial tree node.
struct BinomialTree<T> {
    value: T,
    order: usize,
    subtrees: Vec<Box<BinomialTree<T>>>,
}

impl<T> BinomialTree<T> {
    /// Creates a new binomial tree with the given value and order.
    fn new(value: T, order: usize) -> Self {
        BinomialTree {
            value,
            order,
            subtrees: Vec::new(),
        }
    }
}

/// A binomial heap.
pub struct BinomialHeap<T, C: HeapCompare<T>> {
    heap_type: HeapType,
    compare_func: C,
    num_values: usize,
    roots: Vec<Option<Box<BinomialTree<T>>>>,
}

impl<T, C: HeapCompare<T>> BinomialHeap<T, C> {
    /// Creates a new binomial heap with the given heap type and comparison function.
    pub fn new(heap_type: HeapType, compare_func: C) -> Self {
        BinomialHeap {
            heap_type,
            compare_func,
            num_values: 0,
            roots: Vec::new(),
        }
    }

    /// Compares two values based on the heap type.
    fn compare(&self, a: &T, b: &T) -> Ordering {
        match self.heap_type {
            HeapType::Min => self.compare_func.compare(a, b),
            HeapType::Max => self.compare_func.compare(b, a),
        }
    }

    /// Merges two binomial trees.
    fn merge_trees(&self, mut tree1: Box<BinomialTree<T>>, mut tree2: Box<BinomialTree<T>>) -> Box<BinomialTree<T>> {
        if self.compare(&tree1.value, &tree2.value) == Ordering::Greater {
            mem::swap(&mut tree1, &mut tree2);
        }

        let mut new_tree = Box::new(BinomialTree::new(tree1.value, tree1.order + 1));
        new_tree.subtrees.extend(tree1.subtrees.into_iter());
        new_tree.subtrees.push(tree2);
        new_tree
    }

    /// Merges two binomial heaps.
    fn merge(&mut self, other: &mut Self) {
        let mut carry = None;
        let max_length = self.roots.len().max(other.roots.len());

        for i in 0..max_length {
            let mut vals = [
                self.roots.get_mut(i).and_then(|r| r.take()),
                other.roots.get_mut(i).and_then(|r| r.take()),
                carry.take(),
            ];

            let (new_root, new_carry) = match vals {
                [Some(a), Some(b), Some(c)] => {
                    let merged = self.merge_trees(a, b);
                    (Some(c), Some(merged))
                }
                [Some(a), Some(b), None] => {
                    let merged = self.merge_trees(a, b);
                    (None, Some(merged))
                }
                [Some(a), None, Some(c)] => {
                    let merged = self.merge_trees(a, c);
                    (None, Some(merged))
                }
                [None, Some(b), Some(c)] => {
                    let merged = self.merge_trees(b, c);
                    (None, Some(merged))
                }
                [Some(a), None, None] => (Some(a), None),
                [None, Some(b), None] => (Some(b), None),
                [None, None, Some(c)] => (Some(c), None),
                [None, None, None] => (None, None),
            };

            if new_root.is_some() {
                if i >= self.roots.len() {
                    self.roots.push(new_root);
                } else {
                    self.roots[i] = new_root;
                }
            }

            carry = new_carry;
        }

        if carry.is_some() {
            self.roots.push(carry);
        }

        self.num_values += other.num_values;
        other.roots.clear();
        other.num_values = 0;
    }

    /// Inserts a value into the heap.
    pub fn insert(&mut self, value: T) {
        let new_tree = Box::new(BinomialTree::new(value, 0));
        let mut fake_heap = BinomialHeap::new(self.heap_type, self.compare_func.clone());
        fake_heap.roots.push(Some(new_tree));
        fake_heap.num_values = 1;

        self.merge(&mut fake_heap);
    }

    /// Pops the smallest (or largest) value from the heap.
    pub fn pop(&mut self) -> Option<T> {
        if self.num_values == 0 {
            return None;
        }

        let (least_index, _) = self.roots
            .iter()
            .enumerate()
            .filter_map(|(i, r)| r.as_ref().map(|t| (i, &t.value)))
            .min_by(|(_, a), (_, b)| self.compare(a, b))?;

        let mut least_tree = self.roots[least_index].take().unwrap();
        self.roots[least_index] = None;

        let mut fake_heap = BinomialHeap::new(self.heap_type, self.compare_func.clone());
        fake_heap.roots = least_tree.subtrees.drain(..).map(Some).collect();
        fake_heap.num_values = least_tree.order;

        self.merge(&mut fake_heap);

        self.num_values -= 1;
        Some(least_tree.value)
    }

    /// Returns the number of entries in the heap.
    pub fn num_entries(&self) -> usize {
        self.num_values
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct IntCompare;

    impl HeapCompare<i32> for IntCompare {
        fn compare(&self, a: &i32, b: &i32) -> Ordering {
            a.cmp(b)
        }
    }

    #[test]
    fn test_insert_and_pop() {
        let mut heap = BinomialHeap::new(HeapType::Min, IntCompare);
        heap.insert(3);
        heap.insert(1);
        heap.insert(4);
        heap.insert(2);

        assert_eq!(heap.pop(), Some(1));
        assert_eq!(heap.pop(), Some(2));
        assert_eq!(heap.pop(), Some(3));
        assert_eq!(heap.pop(), Some(4));
        assert_eq!(heap.pop(), None);
    }

    #[test]
    fn test_num_entries() {
        let mut heap = BinomialHeap::new(HeapType::Min, IntCompare);
        assert_eq!(heap.num_entries(), 0);

        heap.insert(1);
        assert_eq!(heap.num_entries(), 1);

        heap.insert(2);
        assert_eq!(heap.num_entries(), 2);

        heap.pop();
        assert_eq!(heap.num_entries(), 1);

        heap.pop();
        assert_eq!(heap.num_entries(), 0);
    }
}
